

class CostGraph:
    '''
    A graph data structure to simplify the edge cost graph. It has two main functions:
    1. To feed the quadratic resharding costs into solver, we need to linearize it. We build edge_cost in
    CostGraph, and it stored every combinations of strategies for a src-dst node pair in an 1D list.
    2. To reduce the searching space, we merge computationally-trivial operators, such as
    element-wise operators, transpose, and reduction, into their following nodes. The merging information will
    be given by the StrategiesVector depending on the type of target node and following nodes.

    Argument:
        leaf_strategies(List[StrategiesVector]): It stores StrategiesVector of every nodes on the graph.
        simplify(bool, optional): The generated cost graph will be simplified if it is true. (default to True)
    '''

    def __init__(self, leaf_strategies, strategy_map, simplify=True, forward_only=True):
        self.leaf_strategies = leaf_strategies
        self.strategy_map =strategy_map
        self.nodes = [strategies_vector.node for strategies_vector in self.leaf_strategies]
        self.nodes_name = [strategies_vector.node.name for strategies_vector in self.leaf_strategies]
        # stores number of strategies in each node
        self.node_lens = {strategies_vector.node: len(strategies_vector) for strategies_vector in self.leaf_strategies}
        # extra_node_costs will store the extra costs introduced by merging nodes
        self.extra_node_costs = {}
        self.following_dict = {}
        self.simplify = simplify
        self.forward_only = forward_only
        self.parent_nodes={}
        self.children_nodes={}
        self._build_cost_graph()


    def _build_cost_graph(self):
        '''
        This method will generate edge_cost for adjacent node pair. Additionally, 'parents' and 'children' attribute will be
        set to node.
        '''
        self.edge_costs = {}
        if self.simplify:
            self.merge_pair = []
        for strategies_vector in self.leaf_strategies:
            # build edge_cost
            dst_node = strategies_vector.node
            for src_node in strategies_vector.predecessor_nodes:
                if src_node not in self.nodes_name:
                    continue
                node_pair = (src_node, dst_node.name)
                # print(node_pair)
                edge_cost = {}
                for i in range(len(strategies_vector)):
                    # node没有strategies_vector属性从map得到
                    for j in range(len(self.strategy_map[src_node])):
                        redistribute_cost_item = strategies_vector[i].redistribute_costs[src_node][j]
                        if self.forward_only:
                            edge_cost[(j, i)] = redistribute_cost_item
                        # else:
                        #     edge_cost[(j, i)] = redistribute_cost_item.total
                # print(edge_cost)
                    self.edge_costs[node_pair] = edge_cost
                    self.parent_nodes[dst_node.name] = []   #TODO 有什么用这两个参数
                    self.children_nodes[dst_node.name] = []
                    for node in strategies_vector.predecessor_nodes:
                        self.parent_nodes[dst_node.name].append(node)
                    for node in strategies_vector.successor_nodes:
                        self.children_nodes[dst_node.name].append(node)
                    for followed_node in strategies_vector.predecessor_nodes:
                        self.merge_pair.append((followed_node, dst_node.name))
            # if self.simplify and strategies_vector.check_merge():
            #     for followed_node in strategies_vector.predecessor_nodes:
            #         # we only merge node pairs which src node has a tensor element inside.
            #         # This is necessary because the node without a tensor element inside will not
            #         # be assigned any strategy.
            #         if _check_tensor_in_node(followed_node._meta_data):
            #             self.merge_pair.append((followed_node, dst_node))

    def simplify_graph(self):
        print("merge_pair:  ",self.merge_pair)
        print("self.children_nodes: ",self.children_nodes)
        print("self.parent_nodes: ",self.parent_nodes)
        self.merge_pair.reverse()
        for (src_node, dst_node) in self.merge_pair:
            self.merge_node(src_node, dst_node)
        self.merge_pair.reverse()
        reindexing_following_dict = {}
        print("self.following_dict: ",self.following_dict)
        for dst, src in self.following_dict.items():
            reindexing_following_dict[dst] = self._reindexing_src(src)
        print(reindexing_following_dict)
        self.following_dict = reindexing_following_dict

    def merge_node(self, src_node, dst_node):
        # print("merge_node")
        # print(self.parent_nodes[dst_node])
        for child_node in self.children_nodes[dst_node]:
            if len(self.parent_nodes[dst_node]) != 0:
                self.following_dict[dst_node] = src_node
                self.children_nodes[dst_node] = []
    def _reindexing_src(self, src):
        if src not in self.following_dict:
            return src
        return self._reindexing_src(self.following_dict[src])